import uuid
from io import BytesIO

from fastapi import FastAPI, HTTPException
from langchain_core.messages import HumanMessage
from langgraph.constants import END
from langgraph.graph import add_messages
from starlette.middleware.cors import CORSMiddleware
import uvicorn
#from app.functions.classify_issue_or_require import process_initial_input_node
from app.graph import DataRequirementAnalysisWorkflow, DataRequirementAnalysisWorkflow2
from matplotlib import image as mpimg, pyplot as plt
import requests

from app.models import SessionState, ChatRequest, ChatResponse

# 初始化FastAPI应用
app = FastAPI(title="数据需求分析工作流API")

# 允许跨域请求
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# 初始化工作流
workflow = DataRequirementAnalysisWorkflow2().graph

@app.post("/chat", response_model=ChatResponse)
def chat(request: ChatRequest):
    # 处理线程ID - 存在则使用，不存在则创建新的
    if request.thread_id:
        thread_id = request.thread_id
        config = {"configurable": {"thread_id": thread_id}}

        # 检查线程是否存在
        try:
            # 尝试获取状态来验证线程是否存在
            workflow.get_state(config)
        except Exception:
            raise HTTPException(status_code=404, detail="线程ID不存在")
    else:
        # 创建新线程
        thread_id = f"user_thread_{uuid.uuid4().hex[:8]}"
        config = {"configurable": {"thread_id": thread_id}}

        # 初始化状态
        initial_state = {
            "original_input": None,
            "input_type": None,
            "require_clearness": None,
            "clarify_input": None,
            "require_json": None,
            "field_availability": None,
            "field_json": None,
            "process_solution": None,
            "collect_solution": None,
            "change_or_not": None,
            "next": None,
            "messages": [],
            "user_initial_input": None,
            "missing_fields": None,
            "user_require": [],
            "user_loop_count": 0,
            "user_satisfaction_intention": 0,
            "user_update_fileds": None
        }
        workflow.update_state(config, initial_state)
         # 关键修复：新会话创建后立即执行到第一个断点
        # 这会触发从START -> human_initial -> classify_user_input的流程
        for event in workflow.stream(None, config, stream_mode="values"):
            initial_state = event
        workflow.update_state(config, initial_state)

    # 获取当前状态
    current_state = workflow.get_state(config).values
    # 添加用户输入到消息列表
    if request.user_input:
        current_state["messages"] = add_messages(
            current_state["messages"],
            [HumanMessage(content=request.user_input)]
        )
        workflow.update_state(config, current_state)
     # 运行流程直到断点
    for event in workflow.stream(None, config, stream_mode="values"):
        current_state = event

    # 检查是否结束
    is_complete = current_state.get("next") == END

    # 获取助手消息
    assistant_response = current_state["messages"][-1].content if current_state.get("messages") else ""

    return {
        "thread_id": thread_id,
        "assistant_response": assistant_response,
        "is_complete": is_complete
    }

# if __name__ == "__main__":
#     workflow = DataRequirementAnalysisWorkflow2().graph
#
#     thread_id = "user_thread_001"
#     config = {"configurable": {"thread_id": thread_id}}
#
#
#     state: SessionState={
#         "original_input": None,
#         "input_type": None,
#         "require_clearness": None,
#         "clarify_input": None,
#         "require_json": None,
#         "field_availability": None,
#         "field_json": None,
#         "process_solution": None,
#         "collect_solution": None,
#         "change_or_not": None,
#         "next": None,  # 初始不指定下一步，由流程逻辑决定
#         "messages": [] ,# 消息列表初始化为空列表
#         "user_initial_input": None ,  #
#         "missing_fields": None,  # 缺失字段信息
#         "user_require":[],
#         "user_loop_count":0, #用户多次更新字段需求最大次数，暂定最大3次
#         "user_satisfaction_intention":0, #用户满意程度意图标识，0为不满意，1为满意
#         "user_update_fileds":None #用于存储提取出来的字段信息
#     }
#
#         # 初始化状态
#     workflow.update_state(config, state)
#     # 运行流程直到断点
#     current_state = state
#     while True:
#         # 运行流程直到断点
#         for event in workflow.stream(None, config, stream_mode="values"):
#             current_state = event
#
#         # 显示助手消息
#         if current_state.get("messages"):
#             print(f"助手: {current_state['messages'][-1].content}")
#
#         # 检查是否结束
#         if current_state.get("next") == END:
#             print("会话结束。")
#             break
#
#         # 断点时的输入
#         user_input = input("请输入您的内容: ")
#         current_state["messages"]=add_messages(current_state["messages"], [HumanMessage(content=user_input)])
#         #workflow.update_state(config, {"message":add_messages(current_state["messages"], [HumanMessage(content=user_input)])})
#         workflow.update_state(config, current_state)


if __name__ == "__main__":

    uvicorn.run(app, host="0.0.0.0", port=8000)